{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%run main.py\n",
    "%load_ext autoreload\n",
    "%autoreload 2\n",
    "\n",
    "!mkdir -p {DATA_DIR} {NAVEC_DIR} {MODEL_DIR}\n",
    "s3 = S3()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "metadata": {},
   "outputs": [],
   "source": [
    "if not exists(NERUS):\n",
    "    s3.download(S3_NERUS, NERUS)\n",
    "    \n",
    "if not exists(NAVEC):\n",
    "    !wget {NAVEC_URL} -O {NAVEC}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "metadata": {},
   "outputs": [],
   "source": [
    "navec = Navec.load(NAVEC)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "metadata": {},
   "outputs": [],
   "source": [
    "words_vocab = Vocab(navec.vocab.words)\n",
    "shapes_vocab = Vocab([PAD] + SHAPES)\n",
    "tags_vocab = BIOTagsVocab(TAGS)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.manual_seed(SEED)\n",
    "seed(SEED)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "metadata": {},
   "outputs": [],
   "source": [
    "word = NavecEmbedding(navec)\n",
    "shape = Embedding(\n",
    "    vocab_size=len(shapes_vocab),\n",
    "    dim=SHAPE_DIM,\n",
    "    pad_id=shapes_vocab.pad_id\n",
    ")\n",
    "emb = TagEmbedding(word, shape)\n",
    "encoder = TagEncoder(\n",
    "    input_dim=emb.dim,\n",
    "    layer_dims=LAYER_DIMS,\n",
    "    kernel_size=KERNEL_SIZE,\n",
    ")\n",
    "ner = NERHead(encoder.dim, len(tags_vocab))\n",
    "model = NER(emb, encoder, ner)\n",
    "\n",
    "model = model.to(DEVICE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "application/vnd.jupyter.widget-view+json": {
       "model_id": "52fcdda9d47d4a4e884564ced860ff51",
       "version_major": 2,
       "version_minor": 0
      },
      "text/plain": [
       "0it [00:00, ?it/s]"
      ]
     },
     "metadata": {},
     "output_type": "display_data"
    }
   ],
   "source": [
    "lines = load_gz_lines(CUSTOM_TEXTS) if CUSTOM_TUNING else load_gz_lines(NERUS)\n",
    "lines = log_progress(lines)\n",
    "\n",
    "items = parse_jl(lines)\n",
    "\n",
    "markups = (SpanMarkup.from_json(_) for _ in items)\n",
    "markups = (_.to_bio(list(tokenize(_.text))) for _ in markups)\n",
    "\n",
    "encode = TagTrainEncoder(\n",
    "    words_vocab, shapes_vocab, tags_vocab,\n",
    "    seq_len=256,\n",
    "    batch_size=64,\n",
    "    shuffle_size=1000,\n",
    ")\n",
    "batches = encode(markups)\n",
    "# full lenta is ~10k batches\n",
    "# epoch takes ~1h, on 2nd epoch loss expodes\n",
    "# ~5k seems to saturate acc\n",
    "batches = head(batches, 5000)\n",
    "batches = [_.to(DEVICE) for _ in batches]\n",
    "\n",
    "size = 10  # ~750 articles\n",
    "batches = {\n",
    "    TEST: batches[:size],\n",
    "    TRAIN: batches[size:]\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "board = MultiBoard([\n",
    "    TensorBoard(BOARD_NAME, RUNS_DIR),\n",
    "    LogBoard()\n",
    "])\n",
    "boards = {\n",
    "    TRAIN: board.section(TRAIN_BOARD),\n",
    "    TEST: board.section(TEST_BOARD),\n",
    "}"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [],
   "source": [
    "optimizer = optim.Adam(model.parameters(), lr=LR)\n",
    "scheduler = optim.lr_scheduler.ExponentialLR(optimizer, LR_GAMMA)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "meters = {\n",
    "    TRAIN: NERScoreMeter(),\n",
    "    TEST: NERScoreMeter(),\n",
    "}\n",
    "\n",
    "for epoch in log_progress(range(EPOCHS)):\n",
    "    model.train()\n",
    "    for batch in log_progress(batches[TRAIN], leave=False):\n",
    "        optimizer.zero_grad()\n",
    "        batch = process_batch(model, ner.crf, batch)\n",
    "        batch.loss.backward()\n",
    "        optimizer.step()\n",
    "    \n",
    "        score = NERBatchScore(batch.loss)\n",
    "        meters[TRAIN].add(score)\n",
    "\n",
    "    meters[TRAIN].write(boards[TRAIN])\n",
    "    meters[TRAIN].reset()\n",
    "\n",
    "    model.eval()\n",
    "    with torch.no_grad():\n",
    "        for batch in log_progress(batches[TEST], leave=False, desc=TEST):\n",
    "            batch = process_batch(model, ner.crf, batch)\n",
    "            batch.pred = ner.crf.decode(batch.pred)\n",
    "            score = score_ner_batch(batch, tags_vocab)\n",
    "            meters[TEST].add(score)\n",
    "\n",
    "        meters[TEST].write(boards[TEST])\n",
    "        meters[TEST].reset()\n",
    "    \n",
    "    scheduler.step()\n",
    "    board.step()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# [2020-04-11 09:35:47]    0 7.3741 01_train/01_loss\n",
    "# [2020-04-11 09:36:09]    0 4.2063 02_test/01_loss\n",
    "# [2020-04-11 09:36:09]    0 0.9776 02_test/02_PER\n",
    "# [2020-04-11 09:36:09]    0 0.9691 02_test/03_LOC\n",
    "# [2020-04-11 09:36:09]    0 0.9486 02_test/04_ORG\n",
    "# [2020-04-11 10:09:54]    1 4.5356 01_train/01_loss\n",
    "# [2020-04-11 10:10:18]    1 3.8757 02_test/01_loss\n",
    "# [2020-04-11 10:10:18]    1 0.9791 02_test/02_PER\n",
    "# [2020-04-11 10:10:18]    1 0.9711 02_test/03_LOC\n",
    "# [2020-04-11 10:10:18]    1 0.9494 02_test/04_ORG\n",
    "# [2020-04-11 10:44:08]    2 4.1505 01_train/01_loss\n",
    "\n",
    "# [2020-04-11 10:44:32]    2 3.7297 02_test/01_loss\n",
    "# [2020-04-11 10:44:32]    2 0.9805 02_test/02_PER\n",
    "# [2020-04-11 10:44:32]    2 0.9729 02_test/03_LOC\n",
    "# [2020-04-11 10:44:32]    2 0.9529 02_test/04_ORG"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "metadata": {},
   "outputs": [],
   "source": [
    "model.emb.shape.dump(MODEL_SHAPE)\n",
    "model.encoder.dump(MODEL_ENCODER)\n",
    "ner.dump(MODEL_NER)\n",
    "\n",
    "# s3.upload(MODEL_SHAPE, S3_MODEL_SHAPE)\n",
    "# s3.upload(MODEL_ENCODER, S3_MODEL_ENCODER)\n",
    "# s3.upload(MODEL_NER, S3_MODEL_NER)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.10.6 ('env': venv)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.6"
  },
  "vscode": {
   "interpreter": {
    "hash": "fbd3b3cf5ce5dbcb71d33f7b8a90c542bb07cc48175c202e830100849640f809"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
